package com.zhou;

import org.apache.pdfbox.io.RandomAccessBuffer;
import org.apache.pdfbox.pdfparser.PDFParser;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.xssf.usermodel.XSSFRow;
import org.apache.poi.xssf.usermodel.XSSFSheet;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

import java.io.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Hello world!
 */
public class App {
    private final String pdfPath = "pdf";
    private final String txtPath = "txt";
    private Set<String> pdfFileNames = new HashSet<>();
    private Set<String> txtFileNames = new HashSet<>();
    private Map<String, List<String>> title2content = new HashMap<>();
    private Pattern pattern = Pattern.compile("([\\s\\S]{0,100}互联网金融[\\s\\S]{0,100})");
    private static String keyword = "互联网金融";
    private static int num = 100;

    public App() {
        new File(pdfPath).mkdirs();
        new File(txtPath).mkdirs();
        pattern = Pattern.compile("([\\s\\S]{0," + num + "}" + keyword + "[\\s\\S]{0," + num + "})");
    }

    public static void main(String[] args) {
        if (args != null && args.length > 1) {
            keyword = args[0];
            num = Integer.parseInt(args[1]);
        }


        while (true) {
            App app = new App();
            System.out.println("根据关键词“" + keyword + "”提取前后" + num + "个字的程序，\n" +
                    "【1】pdf转txt同时提取信息(无txt文件或txt文件不全)\n" +
                    "【2】pdf转txt(仅仅转为txt, 再用【3】提取信息)\n" +
                    "【3】txt直接提取信息(txt文件齐全)\n" +
                    "【4】退出\n" +
                    "请输入数字选择：");
            Scanner scanner = new Scanner(System.in);
            int type = scanner.nextInt();
            if (type == 1) {
                app.startFromPdf(true);
                app.saveExcel();
                System.out.println("\n\n==============================");
                System.out.println("=========【1】完成！==========");
                System.out.println("==============================\n\n");
            } else if (type == 2) {
                app.startFromPdf(false);
                System.out.println("\n\n==============================");
                System.out.println("=========【2】完成！==========");
                System.out.println("==============================\n\n");
            } else if (type == 3) {
                app.startFromTxt();
                app.saveExcel();
                System.out.println("\n\n==============================");
                System.out.println("=========【3】完成！==========");
                System.out.println("==============================\n\n");
            } else if (type == 4) {
                break;
            } else {
                continue;
            }
        }
    }

    /**
     * 保存到Excel
     */
    public void saveExcel() {
        if (title2content.size() == 0) {
            return;
        }

        XSSFWorkbook workbook = new XSSFWorkbook();
        XSSFSheet sheet = workbook.createSheet();
        XSSFRow row = sheet.createRow(0);
        row.createCell(0).setCellValue("名称");
        row.createCell(1).setCellValue("内容");
        String[] keys = title2content.keySet().toArray(new String[1]);
        for (int i = 0; i < keys.length; i++) {
            row = sheet.createRow(i);
            List<String> contents = title2content.get(keys[i]);
            row.createCell(0).setCellValue(keys[i]);
            for (int j = 0; j < contents.size(); j++) {
                row.createCell(j + 1).setCellValue(contents.get(j));
            }
        }

        try {
            FileOutputStream fos = new FileOutputStream("公司公告-" + keyword + ".xlsx");
            workbook.write(fos);
            workbook.close();
            fos.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 开始pdf转换txt
     */
    public void startFromPdf(boolean isExtract) {
        System.out.println(String.format("开始加载pdf文件"));
        openPdfFiles(new File(pdfPath));
        System.out.println(String.format("pdf文件共%d个", pdfFileNames.size()));

        if (pdfFileNames.size() == 0) {
            return;
        }

        System.out.println("\n===============开始转换===============");
        String[] pdfs = pdfFileNames.toArray(new String[1]);
        for (int i = 0; i < pdfs.length; i++) {
            pdf2txt(pdfs[i], isExtract);
            System.out.println(String.format("(%d / %d) 完成：%s", i + 1, pdfs.length, pdfs[i]));
        }
        System.out.println("转换结束");
    }

    /**
     * 开始提取文本
     */
    public void extractTxt(String title, String content) {
        title = title.replace(".pdf", "").replace(".txt", "");
        content = content.replace(" ", "").replace("\r", "").replace("\n", "");


        Matcher matcher = pattern.matcher(content);
        while (matcher.find()) {
            String value = matcher.group();
            List<String> list = title2content.getOrDefault(title, new ArrayList<>());
            list.add(value);
            title2content.put(title, list);
        }
    }

    /**
     * 加载文本文件
     */
    public void startFromTxt() {
        System.out.println(String.format("开始加载txt文件"));
        openTxtFiles(new File(txtPath));
        System.out.println(String.format("txt文件共%d个", txtFileNames.size()));

        if (txtFileNames.size() == 0) {
            return;
        }

        FileInputStream fis = null;
        BufferedReader br = null;

        String[] txts = txtFileNames.toArray(new String[1]);
        for (int i = 0; i < txts.length; i++) {
            try {
                File f = new File(txts[i]);
                fis = new FileInputStream(txts[i]);
                br = new BufferedReader(new InputStreamReader(fis));
                String content = "";
                String temp = null;
                while ((temp = br.readLine()) != null) {
                    content += temp;
                }
                extractTxt(f.getName(), content);
                System.out.println(String.format("(%d / %d) 完成：%s", i + 1, txts.length, txts[i]));
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                try {
                    br.close();
                    fis.close();
                } catch (Exception e) {
                }
            }
        }
    }

    /**
     * 打开所有的txt文件
     */
    public void openTxtFiles(File file) {
        if (file != null && file.isDirectory()) {
            File[] files = file.listFiles();
            for (File f : files) {
                if (f.isDirectory()) {
                    openPdfFiles(f);
                } else if (f.getName().endsWith(".txt")) {
                    txtFileNames.add(f.getAbsolutePath());
                }
            }
        }
    }

    /**
     * 打开所有的pdf文件
     */
    public void openPdfFiles(File file) {
        if (file != null && file.isDirectory()) {
            File[] files = file.listFiles();
            for (File f : files) {
                if (f.isDirectory()) {
                    openPdfFiles(f);
                } else if (f.getName().endsWith(".pdf")) {
                    pdfFileNames.add(f.getAbsolutePath());
                }
            }
        }
    }

    /**
     * pdf转txt
     *
     * @param pdfPath
     */
    public void pdf2txt(String pdfPath, boolean isExtract) {
        if (pdfPath.endsWith(".pdf")) {
            File file = new File(pdfPath);
            String txtName = file.getName().replace(".pdf", ".txt");

            FileInputStream fis = null;
            PDDocument pdDocument = null;
            FileOutputStream fos = null;
            try {
                fis = new FileInputStream(file);
                PDFParser pdfParser = new PDFParser(new RandomAccessBuffer(fis));
                pdfParser.parse();

                pdDocument = pdfParser.getPDDocument();
                PDFTextStripper stripper = new PDFTextStripper();
                String text = stripper.getText(pdDocument);

                fos = new FileOutputStream(txtPath + "/" + txtName);
                fos.write(text.getBytes());
                fos.close();

                if (isExtract) {
                    extractTxt(txtName, text);
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                try {
                    fos.close();
                    pdDocument.close();
                    fis.close();
                } catch (Exception e) {

                }
            }
        }
    }
}
